# 增加 pygame 库引用

import pygame
from typing import Dict, List, Tuple
from algo_struct.algo import rad_limit
import os
# 仿真器入口 用于训练强化学习 跟踪算法用的简化版本仿真器
# 由于没有GUI界面 
# 航控需要通过 位置偏差 和 航向偏差 等指标查看控制算法是否有效 
from dynamics_simulation import motionSim
import numpy as np
import matplotlib.pyplot as plt
import socket
import json
from navigation import gpsAndUtm
from collections import namedtuple, OrderedDict
import traceback
from algo_struct.algo import rad_limit
import pickle, os
import math
import gymnasium as gym

# MAP
MAP_WIDTH, MAP_HEIGHT= 800, 600

# 随机本船ego boat出生区域
init_position = pygame.Rect(123, 59, 124, 101)

# others
# 地形
GROUND = [
    pygame.Rect(62, 246, 354, 62),
    pygame.Rect(392, 210, 408, 49),
]

# 停泊船只
OBSTACLES = [
    pygame.Rect(79, 246, 30, 120),
    pygame.Rect(513, 168, 30, 120),
    pygame.Rect(786, 184, 15, 75),
    pygame.Rect(400, 270, 30, 120),
    pygame.Rect(514, 270, 30, 90),
    pygame.Rect(619, 270, 20, 90),
    pygame.Rect(657, 270, 40, 120),
]

# 运动船只

# 可停泊泊位
POSITIONS = [
    pygame.Rect(79, 405, 30, 120),
    pygame.Rect(672, 168, 30, 120),
    pygame.Rect(488, 390, 30, 120),
    pygame.Rect(598, 390, 30, 120),
    pygame.Rect(748, 390, 30, 120),
]

NUM_RAYS = 64
RAY_LENGTH = 500
RAY_STEP = 10
ROBOT_RADIUS = 5
DT = 0.05
MAX_EPISODE_STEPS = 2000

def rsp_to_force(rsp: int):
                a = 2.30549953e-04
                b = 3.88991253e-01
                c = -2.87920929e+02
                return a*rsp*abs(rsp) + b*rsp + c

class Boat:
    def __init__(self, x: float, y: float, psi: float):
        self.x = x
        self.y = y
        self.psi = psi 
        self.u = 0.0
        self.v = 0.0
        self.dpsi = 0.0
        self.nu = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.acc = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        delta_h = 0.1
        
        # boat运动仿真类输入参数
        my_mmg = motionSim.BoatMotionSim(delta_h=delta_h, name='usv')
        self.x_vec = [self.x, self.y, 0.0, 0.0, 0.0, self.psi]
        self.u_vec = [self.u, self.v, 0.0, 0.0, 0.0, self.dpsi]

    def reset(self, x: float, y: float, psi: float):
        self.x = x
        self.y = y
        self.psi = psi
        self.u = 0.0
        self.v = 0.0
        self.dpsi = 0.0
        self.nu = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.acc = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        self.x_vec = [self.x, self.y, 0.0, 0.0, 0.0, self.psi]
        self.u_vec = [self.u, self.v, 0.0, 0.0, 0.0, self.dpsi]

    def step(self, action: np.ndarray):
        # update state
        rspl = action[0]
        rspr = action[1]
        # rsp to force
        f_left = rsp_to_force(rspl)
        f_right = rsp_to_force(rspr)
        # 计算更新 位置 速度
        self.x_vec, self.u_vec, self.nu, self.acc = my_mmg.run_MMG_double_thrust(x=self.x_vec, u=self.u_vec, f_left=f_left, f_right=f_right)
        self.x_vec[2] = rad_limit(self.x_vec[2]) # clip in [-pi, pi] 

class TopDownRaycastEnv(gym.Env):
    metadata = {"render.modes": ["human"], "render_fps":15}

    def __init__(self, use_render: bool = True, seed: int = None):
        super().__init__()
        self.obstacles = OBSTACLES
        self.target = Boat(x=init_position[0], y=init_position[1], heading=0)


# 解析控制参数
def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    delta = recv_obser['rudl']
    rsp = recv_obser['rspl']

    return delta, rsp

# 打包船的状态返回
def encode_data(trans):
    # 按照标准接口编写
    data = json.dumps(trans, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')

# 计算奖励
def get_reward(delta_psi, dpsi, delta):
    # DDPG高频大舵角 得分高 但是这是执行机构做不到的 因此增加累积舵角惩罚
    score = -((0.1*(delta_psi)**2) + (dpsi**2) + (delta**2))

    print(f'score is {score}')
    return score

# test motion sim
if __name__ == '__main__':
    # 与航行控制算法连接TCP
    LOCAL_HOST = "127.0.0.1"
    LOCAL_PORT = 10906
    
    socket_server = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
    socket_server.bind((LOCAL_HOST, LOCAL_PORT))
    socket_server.listen()
    conn, address = socket_server.accept()
    print(f"收到了客户端的连接，客户端信息是 {address}")
    HOME = os.getcwd()

    path_x = [0, 200, 400, 600, 800]
    path_y = [0, 200, 500, 300, 400]
    # 规划速度：根据距离动态调整速度
    current_p_index = 0  # 当前目标点索引

    boat_data = dict()

    delta_h = 0.1
    
    # boat运动仿真类输入参数
    my_mmg = motionSim.BoatMotionSim(delta_h=delta_h, name='usv')
    
    recorder = dict()
    recorder['exp_anle'] = list()
    recorder['path_x'] = list()
    recorder['path_y'] = list()
    recorder['yaw'] = list()
    recorder['u'] = list()
    recorder['rudder'] = list()
    recorder['y_error'] = list()

    try:
        x = np.array([-50.0, -50.0, 0.0, 0.0, 0.0, 0.0])# 初始无人艇位置
        u = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])# 初始无人艇速度
        current_p_index = 0
        old_rsp = 0
        while (1):
            if (current_p_index > len(path_x) - 1):
                break
            exp_p = (path_x[current_p_index], path_y[current_p_index])
            
            # 动态计算当前目标点的速度
            distance_to_target = np.sqrt((exp_p[0] - x[0]) ** 2 + (exp_p[1] - x[1]) ** 2)
            max_speed = 2.0  # 最大速度
            min_speed = 0.0  # 最小速度
            exp_speed = max(min_speed, min(max_speed, (distance_to_target / 30.0)))  # 根据距离动态调整速度

            if distance_to_target < 30:
                current_p_index += 1
                continue
            
            # 接收转速和舵角
            info, address = conn.recvfrom(1024)
            old_rsp = rsp
            delta, rsp = decode_data(info)

            # 计算更新 位置 速度
            x, u, nu, acc = my_mmg.runMMG(x=x, u=u, delta=delta * np.pi / 180.0, f=rsp)

            # position to exp_angle
            exp_angle = np.arctan2(exp_p[0] - x[0], exp_p[1] - x[1]) * 180.0 / np.pi

            # 为强化学习用
            observation = dict()
            observation['psi'] = rad_limit(rad_limit(x[2]) - exp_angle * np.pi / 180.0) * 180.0 / np.pi  # 角度偏差
            observation['dpsi'] = u[2] * 180.0 / np.pi  # Z轴角速度 弧度转角度
            observation['rud'] = delta  # 舵角 角度
            observation['dspeed'] = exp_speed - u[0]  # x速度偏差 船体坐标系下的轴向速度
            observation['v_x'] = u[0]  # 船体坐标系x速度
            observation['v_y'] = u[1]  # 船体坐标系y速度
            observation['acc_x'] = acc[0]
            observation['acc_n'] = acc[2] * 180.0 / np.pi
            observation['drsp'] = 0.01 * (rsp - old_rsp)  # 转速波动过大惩罚 收敛不需要转速归零
            observation['rsp'] = rsp * 0.01  # 转速

            boat_data['reward'] = get_reward(observation)
            boat_data['observation'] = observation
            boat_data['terminated'] = 0
            boat_data['truncated'] = 0
            data = encode_data(trans=boat_data)

            # 发送数据
            conn.send(data)

            # 计算横向偏差 点到直线距离
            if current_p_index > 0:
                y_error = distance_point_to_line(x[0], x[1], path_x[current_p_index], path_y[current_p_index], path_x[current_p_index - 1],
                                                 path_y[current_p_index - 1])
            else:
                y_error = np.sqrt((exp_p[0] - x[0]) ** 2 + (exp_p[1] - x[1]) ** 2)

            print(f'psi is {x[5]}, exp_angle is {exp_angle}')
            recorder['path_x'].append(x[0])
            recorder['path_y'].append(x[1])
            recorder['yaw'].append(x[5])
            recorder['u'].append(u)
            recorder['exp_anle'].append(exp_angle)
            recorder['rudder'].append(delta)
            recorder['y_error'].append(y_error)

            # 绘图pygame
            

        plt.figure()
        plt.plot(path_x, path_y)
        plt.plot(recorder['path_x'], recorder['path_y'])
        plt.figure()
        plt.plot(recorder['rudder'])
        plt.show()
        
    except Exception:
        # 关闭连接
        conn.close()
        socket_server.close()
        traceback.print_exc()

